{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2238f6dc-4462-4c86-8515-4d5cb6da17e7",
   "metadata": {},
   "source": [
    "# Building Scalable API Backends with Covalent for LLM and Generative AI Applications\n",
    "\n",
    "In this tutorial, we'll delve into the intricacies of constructing scalable API backends for Large Language Models (LLMs) and Generative AI applications. We aim to facilitate seamless collaboration between two cornerstone roles in contemporary machine learning projects: **researchers**, who innovate and experiment with models, and **engineers**, tasked with transforming these models into production-grade applications.\n",
    "\n",
    "Navigating the deployment of high-compute API endpoints, particularly for Generative AI and LLMs, often presents a myriad of challenges. From juggling multiple cloud resources to managing operational overheads and switching between disparate development environments, the endeavor can quickly escalate into a complex ordeal. This tutorial is designed to guide you through these hurdles using [Covalent](https://www.covalent.xyz/), a Pythonic workflow orchestration platform.\n",
    "\n",
    "### Key Challenges and how Covalent solves them\n",
    "- **Resource Management**: The manual management of cloud resources like GPUs is not only tedious but also prone to errors. Covalent automates this, allowing for smooth workflow management right from your Python environment.\n",
    "  \n",
    "- **Operational Overhead**: Tasks like maintaining server uptime, load balancing, and API versioning can complicate the development process. Covalent streamlines these operational aspects, freeing you to focus on development.\n",
    "  \n",
    "- **Environment Switching**: The need to switch between development, testing, and production environments can be a bottleneck, especially in agile, iterative development cycles. Covalent offers a unified environment, simplifying this transition.\n",
    "  \n",
    "- **Multi-Cloud Deployment**: With GPUs often in short supply, the ability to deploy across multiple cloud providers is increasingly crucial. Covalent supports multi-cloud orchestration, making this usually complex task straightforward.\n",
    "  \n",
    "- **Scalability**: High-compute tasks often require dynamic scaling, which can be cumbersome to manage manually. Covalent handles this automatically, adapting to the computational needs of your project.\n",
    "\n",
    "### Tutorial overview\n",
    "\n",
    "This tutorial will encompass the following steps:\n",
    "\n",
    "1. Developing a customizable Covalent [workflow designed to employ AI for news article summarization](#News-Summarization-workflow) [***researcher***],\n",
    "2. [Executing experiments on the established Covalent workflows iteratively](#Rerunning-Workflows), aiming for desirable performance outcomes [***researcher***], and\n",
    "3. [Rerunning and reusing experiments via the Streamlit application](Rerunning-workflows-via-Streamlit) [***engineer***]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b689ccd-c0b8-47cb-a953-489f41587075",
   "metadata": {},
   "source": [
    "# Getting started\n",
    "\n",
    "This tutorial requires [PyTorch](https://pytorch.org/), [Diffusers](https://github.com/huggingface/diffusers), [Hugging Face Transformers](https://huggingface.co/docs/transformers/index) for generative AI. [Streamlit](https://streamlit.io/) will serve to make the user experience smooth. To install all of them, simply use the `requirements.txt` file to replicate this notebook. \n",
    "\n",
    "The list of packages required to run this tutorial is listed below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "03a0c441-40ec-4353-82eb-4d395fb8e4e2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "accelerate==0.21.0\n",
      "bs4==0.0.1\n",
      "covalent-azurebatch-plugin==0.12.0\n",
      "diffusers==0.19.3\n",
      "emoji==2.8.0\n",
      "Pillow==9.5.0\n",
      "sentencepiece==0.1.99\n",
      "streamlit==1.25.0\n",
      "torch==2.0.1\n",
      "transformers==4.31.0\n",
      "xformers==0.0.21\n"
     ]
    }
   ],
   "source": [
    "with open(\"./requirements.txt\", \"r\") as file:\n",
    "    for line in file:\n",
    "        print(line.rstrip())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "3846e49b-f0fa-4878-b733-5fa53c1452af",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment below line to install necessary libraries\n",
    "# !pip install requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "063f6e62-402f-4232-b6e2-f94568aee059",
   "metadata": {},
   "outputs": [],
   "source": [
    "# save under workflow.py\n",
    "import os\n",
    "import re\n",
    "import requests\n",
    "from uuid import uuid4\n",
    "from bs4 import BeautifulSoup\n",
    "\n",
    "import transformers\n",
    "from transformers import (\n",
    "    AutoTokenizer, T5Tokenizer, T5ForConditionalGeneration,\n",
    "    pipeline, AutoModelForSequenceClassification\n",
    ")\n",
    "from diffusers import DiffusionPipeline\n",
    "from PIL import Image, ImageDraw, ImageFont\n",
    "import covalent as ct\n",
    "import torch\n",
    "import logging"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "957bdb21-b577-4d56-99a5-b07cea92f1dc",
   "metadata": {},
   "source": [
    "# News Summarization workflow\n",
    "\n",
    "\n",
    "We first define executors to use [Azure Batch](https://github.com/AgnostiqHQ/covalent-azurebatch-plugin) as compute. Two types of executors allow us to leverage different executors for different compute "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "13658ad0-c002-4ef6-bc4d-ad9fed565443",
   "metadata": {},
   "outputs": [],
   "source": [
    "# save under workflow.py\n",
    "\n",
    "# setting loggers to info to avoid too many debug messages\n",
    "loggers = [logging.getLogger(name) for name in logging.root.manager.loggerDict]\n",
    "for logger in loggers:\n",
    "    logger.setLevel(logging.INFO)\n",
    "\n",
    "\n",
    "\n",
    "# define dependencies to install on remote execution\n",
    "DEPS_ALL = ct.DepsPip(\n",
    "    packages=[\n",
    "        \"transformers==4.31.0\", \"diffusers==0.19.3\", \"accelerate==0.21.0\",\n",
    "        \"cloudpickle==2.2.0\", \"sentencepiece==0.1.99\", \"torch==2.0.1\",\n",
    "        \"Pillow==9.5.0\", \"xformers==0.0.21\", \"emoji==2.8.0\", \"protobuf\"\n",
    "    ]\n",
    ")\n",
    "azure_cpu_executor = ct.executor.AzureBatchExecutor(\n",
    "    # Ensure to specify your own Azure resource information\n",
    "    pool_id=\"covalent-cpu\",\n",
    "    retries=3,\n",
    "    time_limit=600,\n",
    ")\n",
    "\n",
    "# base_image_uri points to a non-default different docker image to support use nvidia gpu\n",
    "azure_gpu_executor = ct.executor.AzureBatchExecutor(\n",
    "    # Ensure to specify your own Azure resource information\n",
    "    pool_id=\"covalent-gpu\",\n",
    "    retries=3,\n",
    "    time_limit=600,\n",
    "    base_image_uri=\"docker.io/filipbolt/covalent_azure:0.220.0\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a8c2d87-9d9a-4b56-ae9b-8c7b8e2668d2",
   "metadata": {},
   "source": [
    "Each electron is associated with an executor, where the computation takes place. Within this framework, less demanding tasks are allocated to the `cpu` executor, while computationally intensive tasks, like generating images from textual prompts, are designated to the `gpu` for compute resources. First, we provide the task outlines. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "113c6ecd-719a-4be1-940b-d38fefe26733",
   "metadata": {},
   "outputs": [],
   "source": [
    "# save under workflow.py\n",
    "@ct.electron(executor=azure_cpu_executor)\n",
    "def extract_news_content(news_url):\n",
    "    response = requests.get(news_url)\n",
    "    soup = BeautifulSoup(response.content, \"html.parser\")\n",
    "\n",
    "    # Extracting article text\n",
    "    paragraphs = soup.find_all(\"p\")\n",
    "    article = \" \".join([p.get_text() for p in paragraphs])\n",
    "    return article\n",
    "\n",
    "@ct.electron(executor=azure_cpu_executor)\n",
    "def generate_title(\n",
    "    article, model_name=\"JulesBelveze/t5-small-headline-generator\",\n",
    "    max_tokens=84, temperature=1, no_repeat_ngram_size=2\n",
    "):\n",
    "    ...\n",
    "\n",
    "@ct.electron(executor=azure_gpu_executor)\n",
    "def generate_reduced_summary(\n",
    "    article, model_name=\"t5-small\", max_length=30\n",
    "):\n",
    "    ...\n",
    "\n",
    "@ct.electron(executor=azure_cpu_executor)\n",
    "def add_title_to_image(image, title):\n",
    "    ...\n",
    "\n",
    "@ct.electron(executor=azure_gpu_executor)\n",
    "def sentiment_analysis(\n",
    "    article, model_name=\"finiteautomata/bertweet-base-sentiment-analysis\"\n",
    "):\n",
    "    ...\n",
    "\n",
    "@ct.electron(executor=azure_cpu_executor)\n",
    "def generate_image_from_text(\n",
    "    reduced_summary, model_name=\"OFA-Sys/small-stable-diffusion-v0\", prompt=\"Impressionist image - \"\n",
    "):\n",
    "    ...\n",
    "\n",
    "@ct.electron(executor=azure_cpu_executor)\n",
    "def save_image(image, filename='image'):\n",
    "    ..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d9f45c3-71ff-4e2f-83ae-6b45bbde67ee",
   "metadata": {},
   "source": [
    "## Covalent Workflow\n",
    "\n",
    "The workflow connects all these steps (electrons) into a workflow (lattice) into a cohesive and runnable workflow. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "8dde7bfe-ead8-4b4b-b2c8-3633f0af9c29",
   "metadata": {},
   "outputs": [],
   "source": [
    "# save under workflow.py\n",
    "@ct.lattice\n",
    "def news_content_curator(\n",
    "    news_url, image_generation_prefix=\"Impressionist image \",\n",
    "    summarizer_model=\"t5-small\",\n",
    "    summarizer_max_length=40,\n",
    "    title_generating_model=\"JulesBelveze/t5-small-headline-generator\",\n",
    "    image_generation_model=\"OFA-Sys/small-stable-diffusion-v0\",\n",
    "    temperature=1, max_tokens=64, no_repeat_ngram_size=2,\n",
    "    content_analysis_model=\"finiteautomata/bertweet-base-sentiment-analysis\"\n",
    "):\n",
    "    article = extract_news_content(news_url)\n",
    "    content_property = sentiment_analysis(\n",
    "        article, model_name=content_analysis_model\n",
    "    )\n",
    "    reduced_summary = generate_reduced_summary(\n",
    "        article, model_name=summarizer_model, max_length=summarizer_max_length\n",
    "    )\n",
    "    title = generate_title(\n",
    "        article, model_name=title_generating_model,\n",
    "        temperature=temperature, max_tokens=max_tokens,\n",
    "        no_repeat_ngram_size=no_repeat_ngram_size\n",
    "    )\n",
    "    generated_image = generate_image_from_text(\n",
    "        reduced_summary, prompt=image_generation_prefix,\n",
    "        model_name=image_generation_model\n",
    "    )\n",
    "    image_with_title = add_title_to_image(generated_image, title)\n",
    "    url = save_image(image_with_title)\n",
    "    return {\n",
    "        \"content_property\": content_property, \"summary\": reduced_summary,\n",
    "        \"title\": title, \"image\": url,\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebc4ec17-82a8-4fab-a6f2-1e77a685d35a",
   "metadata": {},
   "source": [
    "Finally, once a lattice is defined, you must dispatch a workflow to run it. You can dispatch a lattice workflow using Covalent by calling `ct.dispatch` and providing a workflow name and parameters. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "afbe7311-28e8-4f14-8598-e3378ebec946",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "36d7f373-705a-46d6-8ac5-ed57cac8e332\n"
     ]
    }
   ],
   "source": [
    "news_url = 'https://www.quantamagazine.org/math-proof-draws-new-boundaries-around-black-hole-formation-20230816/'\n",
    "dispatch_id = ct.dispatch(news_content_curator)(news_url)\n",
    "print(dispatch_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3288e10f-b1e5-4524-bce8-cfde58b1ae84",
   "metadata": {},
   "source": [
    "The resulting workflow should look like the example below\n",
    "\n",
    "![NewsSum](assets/workflow.gif \"News summarization AI workflow\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83556535-e8af-42ff-958b-79f7233b51f2",
   "metadata": {},
   "source": [
    "Now that the workflow successfully runs, we add more logic to the stub tasks we previously built. \n",
    "\n",
    "Generating text, images, and analyzing content via sentiment analysis can all be implemented via the `transformers` and `diffusers` frameworks: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "83724676-0bb2-4248-a8ce-08fa5781f706",
   "metadata": {},
   "outputs": [],
   "source": [
    "# place in workflow.py\n",
    "@ct.electron(executor=azure_gpu_executor, deps_pip=DEPS_ALL)\n",
    "def generate_title(\n",
    "    article, model_name=\"JulesBelveze/t5-small-headline-generator\",\n",
    "    max_tokens=84, temperature=1, no_repeat_ngram_size=2\n",
    "):\n",
    "    WHITESPACE_HANDLER = lambda k: re.sub(\"\\s+\", \" \", re.sub(\"\\n+\", \" \", k.strip()))\n",
    "\n",
    "    if 't5' in model_name:\n",
    "        tokenizer = T5Tokenizer.from_pretrained(\n",
    "            model_name, legacy=False\n",
    "        )\n",
    "    else:\n",
    "        tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "\n",
    "    model = T5ForConditionalGeneration.from_pretrained(model_name)\n",
    "\n",
    "    # Process and generate title\n",
    "    input_ids = tokenizer(\n",
    "        [WHITESPACE_HANDLER(article)], return_tensors=\"pt\",\n",
    "        padding=\"max_length\", truncation=True, max_length=384,\n",
    "    )[\"input_ids\"]\n",
    "\n",
    "    output_ids = model.generate(\n",
    "        input_ids=input_ids, max_length=max_tokens,\n",
    "        no_repeat_ngram_size=no_repeat_ngram_size, num_beams=4,\n",
    "        temperature=temperature\n",
    "    )[0]\n",
    "\n",
    "    return tokenizer.decode(output_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)\n",
    "    \n",
    "@ct.electron(executor=azure_gpu_executor, deps_pip=DEPS_ALL)\n",
    "def generate_reduced_summary(\n",
    "    article, model_name=\"t5-small\", max_length=30\n",
    "):\n",
    "    if 't5' in model_name:\n",
    "        tokenizer = AutoTokenizer.from_pretrained(model_name + \"_tokenizer\", legacy=False)\n",
    "    else:\n",
    "        tokenizer = T5Tokenizer.from_pretrained(model_name + \"_tokenizer\")\n",
    "\n",
    "    model = T5ForConditionalGeneration.from_pretrained(model_name)\n",
    "\n",
    "    # Encode the article and generate a title\n",
    "    input_text = \"summarize: \" + article\n",
    "    inputs = tokenizer.encode(\n",
    "        input_text, return_tensors=\"pt\", max_length=512, truncation=True\n",
    "    )\n",
    "    # Generate a title with a maximum of max_length words\n",
    "    outputs = model.generate(inputs, max_length=max_length, num_beams=4, length_penalty=2.0, early_stopping=True)\n",
    "    return tokenizer.decode(outputs[0], skip_special_tokens=True)\n",
    "\n",
    "\n",
    "@ct.electron(executor=azure_gpu_executor, deps_pip=DEPS_ALL)\n",
    "def sentiment_analysis(\n",
    "    article,\n",
    "    model_name=\"finiteautomata/bertweet-base-sentiment-analysis\"\n",
    "):\n",
    "    sentiment_pipeline = pipeline(\n",
    "        \"sentiment-analysis\", model=model_name,\n",
    "        padding=True, truncation=True\n",
    "    )\n",
    "    mapping = {\n",
    "        'NEU': 'neutral',\n",
    "        'NEG': 'negative',\n",
    "        'POS': 'positive'\n",
    "    }\n",
    "    label = sentiment_pipeline(article)[0][\"label\"]\n",
    "    return mapping.get(label, label)\n",
    "\n",
    "@ct.electron(executor=azure_gpu_executor, deps_pip=DEPS_ALL)\n",
    "def generate_image_from_text(reduced_summary, model_name=\"OFA-Sys/small-stable-diffusion-v0\", prompt=\"Impressionist image - \"):\n",
    "    model = DiffusionPipeline.from_pretrained(\n",
    "        model_name, safety_checker=None,\n",
    "        torch_dtype=torch.float16\n",
    "    )\n",
    "    model.enable_attention_slicing()\n",
    "    \n",
    "    # Generate image using DiffusionPipeline\n",
    "    reduced_summary = prompt + reduced_summary\n",
    "    _ = model(reduced_summary, num_inference_steps=1)\n",
    "    return model(reduced_summary).images[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8590ad85-6537-49a8-a9f7-ec3b9012be03",
   "metadata": {},
   "source": [
    "The generated images and text can be patched together, and the image may then be uploaded to a cloud storage to make it easier to transfer it via an API. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a010e2cf-b86f-4af1-b513-3ba8753a8d2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "@ct.electron(executor=azure_cpu_executor, deps_pip=DEPS_ALL)\n",
    "def add_title_to_image(image, title):\n",
    "    # Create a new image with space for the title\n",
    "    new_image = Image.new(\n",
    "        \"RGB\", (image.width, image.height + 40), color=\"black\"\n",
    "    )\n",
    "    new_image.paste(image, (0, 40))\n",
    "\n",
    "    # Create a drawing context\n",
    "    draw = ImageDraw.Draw(new_image)\n",
    "    font = ImageFont.load_default()\n",
    "\n",
    "    # Sanitize title to remove non-latin-1 characters\n",
    "    sanitized_title = \"\".join([i if ord(i) < 128 else \" \" for i in title])\n",
    "\n",
    "    # Split the title into multiple lines if it's too long\n",
    "    words = sanitized_title.split()\n",
    "    lines = []\n",
    "    while words:\n",
    "        line = \"\"\n",
    "        while words and font.getlength(line + words[0]) <= image.width:\n",
    "            line += words.pop(0) + \" \"\n",
    "        lines.append(line)\n",
    "\n",
    "    # Calculate position to center the text\n",
    "    y_text = 10\n",
    "    for line in lines:\n",
    "        # Calculate width and height of the text to be drawn\n",
    "        _, _, width, height = draw.textbbox(xy=(0, 0), text=line, font=font)\n",
    "        position = ((new_image.width - width) / 2, y_text)\n",
    "        draw.text(position, line, font=font, fill=\"white\")\n",
    "        y_text += height\n",
    "\n",
    "    return new_image"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a9ab8e9-f61e-4671-961e-91f520e0ae04",
   "metadata": {},
   "source": [
    "Finally, we will upload the save the image to our local machine and transfer it to Streamlit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "bf8c5f33-9b56-45c7-b10e-6fbebb43bfad",
   "metadata": {},
   "outputs": [],
   "source": [
    "@ct.electron\n",
    "def save_image(image, filename='image'):\n",
    "    image.save(f\"{filename}.jpg\")\n",
    "    return image"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "473134d0-c533-45fd-9918-bd92aa9c8de2",
   "metadata": {},
   "source": [
    "# Rerunning Workflows\n",
    "\n",
    "Upon the execution of a Covalent workflow, an associated `dispatch_id` is generated, serving as a unique workflow execution identifier. This dispatch ID serves a dual purpose: it acts as a reference point for the specific workflow and also facilitates the rerun of the entire workflow. Covalent retains a record of all previously executed workflows in a **scalable database**, thus forming a comprehensive history that can be rerun using their respective dispatch IDs.\n",
    "\n",
    "[Redispatching](https://docs.covalent.xyz/docs/features/redispatch/) a workflow to summarize a different news article can be done by providing the `dispatch_id` to the `redispatch` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "264b2e78-8531-45fd-ba77-10cc978756f8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "eb54ecee-1eea-4e96-aca3-a75dd64ba677\n"
     ]
    }
   ],
   "source": [
    "new_url = \"https://www.quantamagazine.org/what-a-contest-of-consciousness-theories-really-proved-20230824/\"\n",
    "redispatch_id = ct.redispatch(dispatch_id)(new_url)\n",
    "print(redispatch_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "922bba8f-8196-49c2-8bd0-473ffdf9af58",
   "metadata": {},
   "source": [
    "It's important to distinguish between dispatching workflows (using `ct.dispatch`) and redispatching them (using `ct.redispatch`). **Dispatching** is typically carried out during the stages of designing a new workflow, while **redispatching** involves replicating and refining a previously created and dispatched workflow."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a159c4fd-a8ac-4e87-8c92-0d62527cc0e9",
   "metadata": {},
   "source": [
    "It's also possible to rerun a workflow while [reusing previously computed results](https://docs.covalent.xyz/docs/features/redispatch/#reuse-previously-computed-results). For instance, if you want to experiment with a different prompt for generating images from the same news article, while keeping the summarization and headline generation unchanged, you can initiate the workflow again, preserving the use of previous results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "62642c3c-2054-403f-a9d5-170bd84f3f02",
   "metadata": {},
   "outputs": [],
   "source": [
    "redispatch_id = ct.redispatch(dispatch_id, reuse_previous_results=True)(new_url, \"Cubistic image\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e296642-6d4d-413b-ab32-00a0782929a0",
   "metadata": {},
   "source": [
    "Furthermore, it's possible to tailor a previously executed workflow by replacing tasks. We can achieve this by employing the `replace_electrons` feature, [which allows us to substitute one task with another](https://docs.covalent.xyz/docs/features/redispatch/#re-executing-the-workflow-with-new-input-arguments). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "e93cdf93-098a-47b7-b464-659fabba20d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "@ct.electron(executor=azure_cpu_executor)\n",
    "def classify_news_genre(\n",
    "    article, model_name=\"abhishek/autonlp-bbc-news-classification-37229289\"\n",
    "):\n",
    "    tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "    model = AutoModelForSequenceClassification.from_pretrained(model_name)\n",
    "\n",
    "    inputs = tokenizer(\n",
    "        article, return_tensors=\"pt\", truncation=True, max_length=512\n",
    "    )\n",
    "    outputs = model(**inputs)\n",
    "    id2label = {\n",
    "        0: \"business\",\n",
    "        1: \"entertainment\",\n",
    "        2: \"politics\",\n",
    "        3: \"sport\",\n",
    "        4: \"tech\"\n",
    "    }\n",
    "    return id2label[outputs.logits.argmax().item()]\n",
    "\n",
    "replace_electrons = {\n",
    "    \"sentiment_analysis\": classify_news_genre\n",
    "}\n",
    "\n",
    "redispatch_id = ct.redispatch(dispatch_id, replace_electrons=replace_electrons)(\n",
    "    new_url, \"Cubistic image\", content_analysis_model=\"abhishek/autonlp-bbc-news-classification-37229289\"\n",
    ")\n",
    "print(redispatch_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05c3f63c-5e34-49d8-a903-5bdc67820c37",
   "metadata": {},
   "source": [
    "# Rerunning workflows via Streamlit\n",
    "\n",
    "\n",
    "Now, let's proceed with the process of constructing the Streamlit app. This app will function as a gateway to Covalent, automatically initializing the Covalent server if it hasn't been started already, and commencing the initial workflow. Subsequently, new workflows will be triggered based on this initial one. \n",
    "\n",
    "At this point, we recommend to decouple the python code into two files:\n",
    "1. `workflow.py` containing the code defining and running the Covalent workflow\n",
    "2. `streamlit_app.py` containing Streamlit code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "46b82422-4acf-4698-bea4-79ea406c5a27",
   "metadata": {},
   "outputs": [],
   "source": [
    "# add to streamlit_app.py\n",
    "import streamlit as st\n",
    "import covalent as ct\n",
    "from subprocess import check_output\n",
    "import subprocess\n",
    "\n",
    "\n",
    "def is_covalent_down():\n",
    "    out = check_output([\"covalent\", \"status\"])\n",
    "    if \"Covalent server is stopped\" in out.decode('utf-8'):\n",
    "        return True\n",
    "    return False\n",
    "\n",
    "\n",
    "def run_covalent_workflow(workflow_filename):\n",
    "    dispatch_id = check_output([\"python\", workflow_filename]).decode(\"utf-8\")\n",
    "    return dispatch_id.strip()\n",
    "\n",
    "\n",
    "def start_covalent():\n",
    "    subprocess.run(\"covalent start --no-cluster\", shell=True)\n",
    "\n",
    "\n",
    "if is_covalent_down():\n",
    "    st.write(\"Covalent is not up. Starting Covalent...\")\n",
    "    start_covalent()\n",
    "    if check_google_creds():\n",
    "        # execute a covalent workflow\n",
    "        dispatch_id = run_covalent_workflow(\"workflow_remote.py\")\n",
    "        # wait for result\n",
    "        ct.get_result(dispatch_id, wait=True)\n",
    "        st.session_state['dispatch_id'] = dispatch_id"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99e29327-2fc4-4be0-aa51-002eda668e13",
   "metadata": {},
   "source": [
    "Now, the Streamlit app will automatically start Covalent server and launch the first workflow. You may also directly then reuse the `dispatch_id` of the launched workflow to start rerunning workflows and iterating with experiments on tweaking the workflow. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "343f3507-bfb8-4afb-9210-72966f861b80",
   "metadata": {},
   "source": [
    "Now that we have the capability to execute and re-execute Covalent workflows, our goal is to offer users a user-friendly interface. Streamlit enables us to achieve precisely that! We have developed a compact Streamlit application that enables users to adjust parameters for the AI news summarization workflow mentioned earlier and trigger previously executed workflows using their dispatch IDs. The sidebar of the Streamlit app will contain the parameters, with some proposed default values, whereas the central part of the Streamlit app will serve to render the results of the Covalent workflows. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f995197b-d196-4dbe-8048-5df6d881168a",
   "metadata": {},
   "source": [
    "The Streamlit sidebar can be defined as:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "0d31b26a-57b1-48fb-bedf-08b710c903d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# add to streamlit_app.py\n",
    "\n",
    "def create_streamlit_sidebar(\n",
    "    stable_diffusion_models, news_summary_generation,\n",
    "    headline_generation_models, sentiment_analysis_models,\n",
    "    genre_analysis_models\n",
    "):\n",
    "    with st.sidebar:\n",
    "        news_article_url = st.text_input(\n",
    "            \"News article URL\",\n",
    "            value=\"https://www.quantamagazine.org/math-proof-draws-new-boundaries-around-black-hole-formation-20230816/\"\n",
    "        )\n",
    "        st.header(\"Parameters\")\n",
    "\n",
    "        # Title generation section\n",
    "        st.subheader(\"Title generation parameters\")\n",
    "        title_generating_model = headline_generation_models[0]\n",
    "        temperature = st.slider(\n",
    "            \"Temperature\", min_value=0.0, max_value=100.0, value=1.0,\n",
    "            step=0.1\n",
    "        )\n",
    "        max_tokens = st.slider(\n",
    "            \"Max tokens\", min_value=2, max_value=50, value=32,\n",
    "        )\n",
    "\n",
    "        # Image generation section\n",
    "        st.subheader(\"Image generation\")\n",
    "        image_generation_prefix = st.text_input(\n",
    "            \"Image generation prompt\",\n",
    "            value=\"impressionist style\"\n",
    "        )\n",
    "        image_generation_model = stable_diffusion_models[0]\n",
    "\n",
    "        # Text summarization section\n",
    "        st.subheader(\"Text summarization\")\n",
    "        summarizer_model = news_summary_generation[0]\n",
    "        summarizer_max_length = st.slider(\n",
    "            \"Summarization text length\", min_value=2, max_value=50, value=20,\n",
    "        )\n",
    "\n",
    "        # Content analysis section\n",
    "        st.subheader(\"Content analysis\")\n",
    "        selected_content_analysis = st.selectbox(\n",
    "            \"Content analysis option\", options=[\n",
    "                \"sentiment analysis\",\n",
    "                \"genre classification\"\n",
    "            ]\n",
    "        )\n",
    "        if selected_content_analysis == \"sentiment analysis\":\n",
    "            content_analysis_model = sentiment_analysis_models[0]\n",
    "        else:\n",
    "            content_analysis_model = genre_analysis_models[0]\n",
    "\n",
    "    return {\n",
    "        'news_url': news_article_url,\n",
    "        'image_generation_prefix': image_generation_prefix,\n",
    "        'summarizer_model': summarizer_model,\n",
    "        'summarizer_max_length': summarizer_max_length,\n",
    "        'title_generating_model': title_generating_model,\n",
    "        'image_generation_model': image_generation_model,\n",
    "        'temperature': temperature,\n",
    "        'max_tokens': max_tokens,\n",
    "        'content_analysis_model': content_analysis_model,\n",
    "        'selected_content_analysis': selected_content_analysis\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60f2975c-9e60-4716-a50e-06a84b8c8941",
   "metadata": {},
   "source": [
    "The central part of the Streamlit app is designed to render results from Covalent server, using the parameters configured in the sidebar. This triggers the generation of an AI-generated summary of the news article, a proposed title, and an AI-generated image depicting the content of the news article."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "f7c9b8e5-23ff-48f3-ad1b-dc7f97eb165c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# add to streamlit_app.py\n",
    "st.title(\"News article AI summarization\")\n",
    "dispatch_id = st.text_input(\"Dispatch ID\")\n",
    "\n",
    "if st.button(\"Generate image and text summary\"):\n",
    "    st.write(\"Generating...\")\n",
    "\n",
    "    container = st.container()\n",
    "\t\t\n",
    "    # select either genre analysis or sentiment analysis\n",
    "    selected_content_analysis = parameters.pop('selected_content_analysis')\n",
    "    if selected_content_analysis != 'sentiment analysis':\n",
    "        replace_electrons = {\n",
    "            \"sentiment_analysis\": ct.electron(classify_news_genre)\n",
    "        }\n",
    "        parameters[\n",
    "            \"content_analysis_model\"\n",
    "        ] = \"abhishek/autonlp-bbc-news-classification-37229289\"\n",
    "    else:\n",
    "        replace_electrons = {}\n",
    "\n",
    "    redispatch_id = ct.redispatch(\n",
    "        dispatch_id, reuse_previous_results=True,\n",
    "        replace_electrons=replace_electrons\n",
    "    )(**parameters)\n",
    "\n",
    "    covalent_info = ct.get_config()['dispatcher']\n",
    "    address = covalent_info['address']\n",
    "    port = covalent_info['port']\n",
    "    covalent_url = f\"{address}:{port}/{redispatch_id}\"\n",
    "\n",
    "    st.write(f\"Covalent URL on remote server: http://{covalent_url}\")\n",
    "\n",
    "    with container:\n",
    "        result = ct.get_result(redispatch_id, wait=True).result\n",
    "        st.subheader(f\"Article generated title: {result['title']}\")\n",
    "        st.write(\n",
    "            \"In terms of \" +\n",
    "            selected_content_analysis +\n",
    "            \" content is: \" + str(result['content_property'])\n",
    "        )\n",
    "        st.image(result['image'])\n",
    "        st.text_area(\n",
    "            label=\"AI generated summary\",\n",
    "            key=\"summary\",\n",
    "            value=result['summary'], disabled=True\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29810af4-a4c1-4bb4-9ff1-375432b4bf7b",
   "metadata": {},
   "source": [
    "If you saved the provided streamlit code in `streamlit_app.py`, you can run it in a separate python console by running\n",
    "\n",
    "`streamlit run streamlit_app.py`\n",
    "\n",
    "This will start the streamlit app on http://localhost:8501\n",
    "\n",
    "You can use the streamlit app as demonstrated below:\n",
    "\n",
    "![StreamlitCovalent](assets/streamlit_workflow.gif \"Streamlit + Covalent\")\n",
    "\n",
    "Generating multiple images with Streamlit via Covalent is demonstrated below\n",
    "\n",
    "![StreamlitCovalent](assets/streamlit_covalent_imagegen.gif \"Streamlit + Covalent\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e721a608-6880-4e1f-a4a1-09018a69798c",
   "metadata": {},
   "source": [
    "# Conclusion\n",
    "\n",
    "You have learned how to build complex machine learning workflows using an example of a news summarization application. A Covalent server takes care of the machine learning workflows, while a Streamlit interface handles user interactions. The two communicate via a single (dispatch) ID, streamlining resource management, enhancing efficiency, and allowing you to concentrate on the machine learning aspects. \n",
    "\n",
    "If you found this interesting, please note that [Covalent](https://github.com/AgnostiqHQ/covalent) is free and [open source](https://www.covalent.xyz/open-source/). Please visit the [Covalent documentation](https://docs.covalent.xyz/docs/) for more information and many more [tutorials](https://docs.covalent.xyz/docs/user-documentation/tutorials/). An example of the Streamlit application described here was deployed [here](https://covalent-news-summary.streamlit.app/). Please note it will not be able to run out of the box, since it requires having valid Azure access credentials. \n",
    "\n",
    "Happy workflow building! 🎈"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
